/*
 * Copyright 2010 Joshua Smith and Nelson Llewellyn
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package norse.plugins.nmap;

import java.io.BufferedReader;
import java.io.File;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.LinkedList;
import java.util.List;
import norse.msg.MessageIF;
import norse.plugins.nmap.error.NmapScanningException;
import norse.plugins.nmap.msg.NmapScanRequest;
import norse.server.AgentIF;
import norse.server.util.CommandLineArg;
import norse.server.model.ScanResults;

/**
 * Class that launches NMAP and returns the XML that NMAP creates.
 */
public class NmapAgent implements AgentIF
{

    /**
     * Constructor.
     * 
     * @param _nmapExecutable the location of the nmap executable.
     */
    public NmapAgent(File _nmapExecutable)
    {
        if (_nmapExecutable == null || !_nmapExecutable.exists()
                || !_nmapExecutable.isFile())
        {
            throw new IllegalArgumentException(
                    "Invalid NMAP location");
        }

        nmapDir = _nmapExecutable.getParentFile();
        nmapCommand = _nmapExecutable.getName();
    }

    /**
     * Runs an NMAP scan and returns the resulting XML as a StringBuilder.
     *
     * @param _config the scan configuration
     * @return the scan results.
     */
    public StringBuilder runNMAPScan(NmapScanConfiguration _config)
            throws NmapScanningException
    {
        Process proc = null;
        try
        {
            List<String> commands = new LinkedList<String>();
            //If we're on a Win7 system, you have to call out cmd.exe first
            if (System.getProperty("os.name").equals("Windows 7"))
            {
                commands.add("cmd.exe");
                commands.add("/c");
            }
            commands.add(nmapCommand);
            _config.addConfigToCommandList(commands);

            ProcessBuilder procBuild = new ProcessBuilder(commands);

            procBuild.directory(nmapDir);
            proc = procBuild.start();

            InputStream stdoutIs = proc.getInputStream();
            InputStreamReader stdoutIsr = new InputStreamReader(stdoutIs);
            BufferedReader stdoutBr = new BufferedReader(stdoutIsr);

            StringBuilder xmlOutput = new StringBuilder();

            InputStream stderrIs = proc.getErrorStream();
            InputStreamReader stderrIsr = new InputStreamReader(stderrIs);
            BufferedReader stderrBr = new BufferedReader(stderrIsr);

            StringBuilder errOutput = new StringBuilder();


//TODO: Not sure if this is the best implementation for this--what happens if there is no output for a while? (It may exit prematurely!)
            String currLine;
            while ((currLine = stdoutBr.readLine()) != null)
            {
                xmlOutput.append(currLine + "\n");
//TODO: Debugging info. REMOVE!
ScanResults.storeResult("<DEBUG><STATUS>NMAP Scan Request received</STATUS><CURR_XML_SIZE>" + xmlOutput.length() + "</CURR_XML_SIZE></DEBUG>");
            }

            while ((currLine = stderrBr.readLine()) != null)
            {
                errOutput.append(currLine + "\n");
            }


            //If there were errors during scanning
//TODO: Add some error handling for less-than-fatal erros.
            if (errOutput.length() != 0)
            {
                //Just throw an exception--the process will be terminated by
                //the finally block below
                throw new NmapScanningException("NMAP error(s): "
                        + errOutput.toString());
            }

            //We've finished with the process and don't want it to have the
            //destroy method called in the finally block, so null it out.
            proc = null;

            return xmlOutput;
        }
        catch (NmapScanningException scanningException)
        {
            //Re-throw the exception
            throw scanningException;
        }
        catch (Exception scanningException)
        {
            throw new NmapScanningException("Exception while scanning: "
                    + scanningException.getMessage(), scanningException);
        }
        finally
        {
            if (proc != null)
            {
                proc.destroy();
            }
        }
    }
    private final File nmapDir;
    private final String nmapCommand;
    private final Class[] myMessages =
    {
        NmapScanRequest.class
    };

    @Override
    public void deliverMessage(MessageIF _message)
    {
        if (!(_message instanceof NmapScanRequest))
        {
            return;
        }

        //TODO: Debugging info. REMOVE!
        ScanResults.storeResult("<DEBUG><STATUS>NMAP Scan Request received</STATUS></DEBUG>");

        try
        {
            List<CommandLineArg> args = new LinkedList<CommandLineArg>();
            args.add(new CommandLineArg("-oX", "-"));
//            //Scans scanme.nmap.org. Good machine to use for examples!
//            args.add(new CommandLineArg("-sS", "scanme.nmap.org"));
            args.add(new CommandLineArg("-sS", "192.168.1.1-254"));
            NmapScanConfiguration config = new NmapScanConfiguration(args);

            //TODO: Debugging info. REMOVE!
            ScanResults.storeResult("<DEBUG><STATUS>Performing NMAP scan</STATUS></DEBUG>");

            StringBuilder xmloutput = runNMAPScan(config);
            ScanResults.storeResult(xmloutput.toString());
        }
        catch (NmapScanningException error)
        {
            //TODO: Debugging info. REMOVE!
            ScanResults.storeResult("<DEBUG><STATUS>Error during NMAP scan: "
                    + error.getMessage() + "</STATUS></DEBUG>");
        }
    }

    public Class[] getMessagesToRegisterWithPostmaster()
    {
        return myMessages;
    }
}
